# data classes for graph construct and storage based on json or former APIs
import json
from dataclasses import dataclass,field

@dataclass
class inst_graph:

    file_name: str=''
    graph_type: str=''
    inst_nodes: list = field(default_factory=list)
    inst_edges: list = field(default_factory=list)

    def get_node_num(self):
        return len(self.inst_nodes)
    
    def get_edge_num(self):
        return len(self.inst_edges)
    
    def to_dict(self):
        return self.__dict__

@dataclass
class inst_node:

    address: str=''
    features: list = field(default_factory=list)
    is_first_inst: bool=False
    bb_id: int=0

    def __init__(self, addr, mnemonic, op_str, is_first_inst, bb_id, max_string_len=40):
        # every node has a feature vector, including 3 features (split asm by ',')
        # each feature is a string, and the length of each string is pre-limited
        # feature[0]: mnemonic, feature[1:2]: op_str
        self.address = addr
        self.features = []
        self.features.append(mnemonic)
        op_str = op_str.split(', ')
        if len(op_str) < 2:
            op_str.append('\x00'*max_string_len)
        for i in range(2):
            if len(op_str[i]) > max_string_len:
                op_str[i] = op_str[i][:max_string_len]
            else:
                op_str[i] += '\x00'*(max_string_len-len(op_str[i]))
        self.features.extend(op_str)
        self.is_first_inst = is_first_inst
        self.bb_id = bb_id
    
    def to_dict(self):
        return self.__dict__

    # @staticmethod
    # def create(addr, mnemonic, op_str, is_first_inst, bb_id, max_string_len=25):
        
    #     self.address = addr
    #     self.features.append(mnemonic)
    #     op_str = op_str.split(', ')
    #     if len(op_str) < 2:
    #         op_str.append('\x00'*max_string_len)
    #     for i in range(2):
    #         if len(op_str[i]) > max_string_len:
    #             op_str[i] = op_str[i][:max_string_len]
    #         else:
    #             op_str[i] += '\x00'*(max_string_len-len(op_str[i]))
    #     self.features.extend(op_str)
    #     self.is_first_inst = is_first_inst
    #     self.bb_id = bb_id

    #     obj=inst_node(addr, mnemonic, op_str, is_first_inst, bb_id, max_string_len)

    #     return obj

@dataclass
class inst_edge:

    src: int=0
    dst: int=0

    def to_dict(self):
        return self.__dict__
    
